package simpledb.execution;

import simpledb.common.DbException;
import simpledb.common.Type;
import simpledb.storage.Tuple;
import simpledb.storage.TupleDesc;
import simpledb.storage.TupleIterator;
import simpledb.transaction.TransactionAbortedException;

import java.util.NoSuchElementException;

import static simpledb.execution.Aggregator.NO_GROUPING;

/**
 * 聚合操作符类，用于计算聚合函数（如 SUM、AVG、MAX、MIN 等）。
 *
 * 注意：我们只支持对单列进行聚合，并且按单列进行分组。
 */
public class Aggregate extends Operator {

    private static final long serialVersionUID = 1L;

    // 子操作符，用于提供输入元组
    private OpIterator        child;

    // 聚合字段在输入元组中的索引（如 SELECT SUM(age) 中的 age 字段）
    private int               agField;

    // 分组字段在输入元组中的索引，-1 表示不分组（对应 SQL 中的 GROUP BY 子句）
    private int               gbField;

    // 分组字段的数据类型（若不分组则为 null）
    private Type              gbFieldType;

    // 聚合操作类型（如 MIN、MAX、SUM、AVG 等）
    private Aggregator.Op     op;

    // 实际执行聚合逻辑的组件（根据字段类型选择 Integer/String 实现）
    private Aggregator        aggregator;

    // 聚合结果的迭代器，用于遍历最终聚合结果
    private TupleIterator     iterator;

    // 输出元组的描述信息（缓存避免重复生成）
    private TupleDesc         td;

    /**
     * 构造函数。
     * 根据聚合字段类型创建对应的聚合器（Integer/String）。
     *
     * @param child 提供输入元组的子操作符（如 SeqScan、Filter 等）
     * @param afield 聚合字段的索引（如 SELECT SUM(age) 中的 age 对应索引）
     * @param gfield 分组字段的索引（-1 表示不分组）
     * @param aop 聚合操作类型
     */
    public Aggregate(OpIterator child, int afield, int gfield, Aggregator.Op aop) {
        this.agField = afield;
        this.gbField = gfield;
        this.op = aop;
        this.child = child;

        TupleDesc originTd = this.child.getTupleDesc();
        this.gbFieldType = (this.gbField == -1 ? null : originTd.getFieldType(this.gbField));
    }

    /**
     * 获取当前聚合操作中所使用的分组字段索引。
     *
     * 如果是不分组聚合，则返回 NO_GROUPING。
     */
    public int groupField() {
        return this.gbField;
    }

    /**
     * 获取分组字段在输入元组中的名称。
     *
     * 如果没有分组字段，则返回 null。
     */
    public String groupFieldName() {
        return this.child.getTupleDesc().getFieldName(this.gbField);
    }

    /**
     * 获取聚合字段的索引。
     */
    public int aggregateField() {
        return this.agField;
    }

    /**
     * 获取聚合字段在输出元组中的名称。
     */
    public String aggregateFieldName() {
        return this.child.getTupleDesc().getFieldName(this.agField);
    }

    /**
     * 获取当前使用的聚合操作。
     */
    public Aggregator.Op aggregateOp() {
        return this.op;
    }

    /**
     * 返回聚合操作的字符串表示。
     */
    public static String nameOfAggregatorOp(Aggregator.Op aop) {
        return aop.toString();
    }

    /**
     * 打开操作符，准备读取数据。
     */
    public void open() throws NoSuchElementException, DbException, TransactionAbortedException {
        super.open();
        this.child.open();

        TupleDesc originTd = this.child.getTupleDesc();

        // 根据聚合字段的类型创建合适的聚合器
        if (originTd.getFieldType(agField) == Type.INT_TYPE) {
            this.aggregator = new IntegerAggregator(this.gbField, this.gbFieldType, this.agField, this.op);
        } else {
            this.aggregator = new StringAggregator(this.gbField, this.gbFieldType, this.agField, this.op);
        }

        // 遍历所有输入元组，进行聚合
        while (this.child.hasNext()) {
            final Tuple tuple = this.child.next();
            this.aggregator.mergeTupleIntoGroup(tuple);
        }

        // 获取聚合后的结果迭代器
        this.iterator = (TupleIterator) this.aggregator.iterator();
        this.iterator.open();
    }

    /**
     * 获取下一个聚合结果元组。
     *
     * 如果有分组字段，则第一个字段是分组字段值，第二个字段是聚合结果；
     * 如果没有分组字段，则只有一个字段代表聚合结果；
     * 如果没有更多结果，则返回 null。
     */
    protected Tuple fetchNext() throws TransactionAbortedException, DbException {
        if (this.iterator.hasNext()) {
            return this.iterator.next();
        }
        return null;
    }

    /**
     * 重置操作符，重新开始读取数据。
     */
    public void rewind() throws DbException, TransactionAbortedException {
        close();
        open();
    }

    /**
     * 关闭操作符，释放资源。
     */
    public void close() {
        this.child.close();
        this.iterator.close();
        super.close();
    }

    /**
     * 获取该聚合操作的输出 TupleDesc。
     *
     * 如果没有分组字段，输出元组只有一个字段（聚合结果）；
     * 如果有分组字段，输出元组有两个字段：
     *   - 第一个是分组字段值；
     *   - 第二个是聚合结果；
     *
     * 聚合字段名应具有可读性，例如："avg(age)"。
     */
    public TupleDesc getTupleDesc() {
        if (this.td != null) {
            return this.td;
        }

        if (this.gbField == NO_GROUPING) {
            Type[] types = new Type[] { Type.INT_TYPE };
            String[] names = new String[] { this.aggregateFieldName() + "(" + this.op + ")" };
            this.td = new TupleDesc(types, names);
        } else {
            Type[] types = new Type[] { this.gbFieldType, Type.INT_TYPE };
            String[] names = new String[] { this.groupFieldName(), this.aggregateFieldName() + "(" + this.op + ")" };
            this.td = new TupleDesc(types, names);
        }

        return this.td;
    }

    @Override
    public OpIterator[] getChildren() {
        return new OpIterator[] { this.child };
    }

    @Override
    public void setChildren(OpIterator[] children) {
        if (children.length > 0) {
            this.child = children[0];
        }
    }
}
